import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from net import SimpleCNN  # 网络模块
import config
import os
from torch.utils.tensorboard import SummaryWriter  # TensorBoard

def train():
    # 数据预处理
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))  # MNIST 均值方差
    ])

    # 加载数据集
    train_dataset = datasets.MNIST(root=config.DATA_ROOT, train=True, transform=transform, download=True)
    test_dataset = datasets.MNIST(root=config.DATA_ROOT, train=False, transform=transform, download=True)

    train_loader = DataLoader(train_dataset, batch_size=config.BATCH_SIZE, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=config.BATCH_SIZE, shuffle=False)

    device = config.DEVICE
    model = SimpleCNN().to(device)
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=config.LEARNING_RATE)

    if not os.path.exists(config.MODEL_SAVE_PATH):
        os.makedirs(config.MODEL_SAVE_PATH)

    # TensorBoard writer
    writer = SummaryWriter(log_dir=os.path.join(config.MODEL_SAVE_PATH, "logs"))

    for epoch in range(config.EPOCHS):
        model.train()
        total_train_loss = 0
        for batch_idx, (images, labels) in enumerate(train_loader):
            images, labels = images.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            total_train_loss += loss.item()

            # 记录每个batch训练loss
            global_step = epoch * len(train_loader) + batch_idx
            writer.add_scalar("train_loss", loss.item(), global_step)

        avg_train_loss = total_train_loss / len(train_loader)

        # 验证集计算val_loss
        model.eval()
        total_val_loss = 0
        with torch.no_grad():
            for images, labels in test_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                loss = criterion(outputs, labels)
                total_val_loss += loss.item()
        avg_val_loss = total_val_loss / len(test_loader)

        # 记录epoch的平均train_loss和val_loss
        writer.add_scalar("epoch_avg_train_loss", avg_train_loss, epoch)
        writer.add_scalar("epoch_avg_val_loss", avg_val_loss, epoch)

        print(f"Epoch [{epoch+1}/{config.EPOCHS}], Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}")

        # 每个epoch保存模型，文件名带epoch号
        torch.save(model.state_dict(), os.path.join(config.MODEL_SAVE_PATH, f"mnist_cnn_epoch_{epoch + 1}.pth"))

    writer.close()
    print("Training finished and model saved.")

if __name__ == "__main__":
    train()
